package cn.six1943.deepchat.module.ai.agent;

import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.scoring.ScoringModel;

import java.util.HashMap;
import java.util.Map;

public class ModelHolder {

    private static Map<Long, ChatModel> genModels = new HashMap<>();
    private static Map<Long, EmbeddingModel> embeddingModels = new HashMap<>();
    private static Map<Long, StreamingChatModel> chatModels = new HashMap<>();
    private static Map<Long, ScoringModel> scoringModels = new HashMap<>();
    private static Map<Long, TokenCountEstimator> estimators = new HashMap<>();

    public static ChatModel getGenModel(long id) {
        ChatModel model = genModels.get(id);
        if (model == null) {
            ModelWithKey mk = ModelCache.getModelWithKey(id);
            model = ModelFactory.createGenModel(mk);
            genModels.put(id, model);
        }
        return model;
    }

    public static EmbeddingModel getEmbeddingModel(long id) {
        EmbeddingModel model = embeddingModels.get(id);
        if (model == null) {
            ModelWithKey mk = ModelCache.getModelWithKey(id);
            model = ModelFactory.createEmbeddingModel(mk);
            embeddingModels.put(id, model);
        }
        return model;
    }

    public static StreamingChatModel getChatModel(long id, String modelName) {
        StreamingChatModel model = chatModels.get(id);
        if (model == null) {
            ModelWithKey mk = ModelCache.getModelWithKey(id);
            mk.setCode(modelName);
            model = ModelFactory.createChatModel(mk);
            chatModels.put(id, model);
        }
        return model;
    }

    public static ScoringModel getScoringModel(long id) {
        ScoringModel model = scoringModels.get(id);
        if (model == null) {
            ModelWithKey mk = ModelCache.getModelWithKey(id);
            model = ModelFactory.createScoringModel(mk);
            scoringModels.put(id, model);
        }
        return model;
    }

    public static TokenCountEstimator getTokenizer(long id) {
        TokenCountEstimator model = estimators.get(id);
        if (model == null) {
            ModelWithKey mk = ModelCache.getModelWithKey(id);
            model = ModelFactory.createTokenizer(mk);
            estimators.put(id, model);
        }
        return model;
    }
}
